import numpy as np
import torch
from torch.utils.data import Dataset


class SignalDataset(Dataset):
    def __init__(self, raw_data):
        # 初始化函数，用于构建对象
        # 参数：
        #   - raw_data: 原始数据，类型为二维数组
        # 属性：
        #   - self._signal: 数据的信号部分，类型为三维浮点型张量 第四列及以后的数据
        #   - self._fea_plus: 数据的特征加强部分，类型为二维浮点型张量 第二列到第四列数据
        #   - self._label: 数据的标签部分，类型为一维长整型张量 第一类数据
        #   - self._emd: 数据的EMD部分，类型为二维浮点型张量（已被注释掉，未使用）
        #  TODO: Fea_PLUS
        Fea_PLUS = 0
        raw_data = raw_data.astype(np.float64)
        self._signal = torch.FloatTensor(raw_data[:, Fea_PLUS + 1:Fea_PLUS + 1 + 268])
        self._fea_plus = torch.FloatTensor(raw_data[:, 1:Fea_PLUS + 1])
        # self._fea_plus1 = torch.FloatTensor(raw_data[:, 1:Fea_PLUS + 1])
        # self._fea_plus2 = torch.FloatTensor(raw_data[:, Fea_PLUS + 1 + 260:Fea_PLUS + Fea_PLUS + 1 + 260])
        # self._fea_plus = np.concatenate((self._fea_plus1, self._fea_plus2), axis=1)
        self._label = torch.LongTensor(raw_data[:, 0])
        # self._label = torch.LongTensor(raw_data[:, -1]).unsqueeze(1)
        # self._emd = torch.FloatTensor(emd_data)

    @property
    def n_insts(self):
        ''' Property for dataset size '''
        return len(self._label)

    # @property
    # def sig_len(self):
    #     # 返回信号的第二维度长度
    #     return self._signal.shape[1]

    def __len__(self):
        return self.n_insts

    def __getitem__(self, idx):
        return self._signal[idx], self._fea_plus[idx], self._label[idx] - 1